import functools

from tqdm import tqdm

from hysdata import HSClient, api
import asyncio
import pandas as pd

from hysdata.store import Store

store = Store()


async def test_speed_limit_async():
    client = HSClient(auto_wait=True, speed_limit_per_minute=120)
    # 异步API调用方式
    df_stock = await client.async_query(api.stock_list())
    codes = df_stock['secu_code'].to_list()
    dfs = []
    for code in tqdm(codes):
        await client.async_auto_wait(api.stock_Info())
        df = await client.async_query(api.stock_Info(en_prod_code=code), cn_column_names=True)
        dfs.append(df)
    return pd.concat(dfs)


def test_speed_limit():
    client = HSClient(auto_wait=True, speed_limit_per_minute=120)
    # 异步API调用方式
    df_stock = client.sync_query(api.stock_list())
    codes = df_stock['secu_code'].to_list()
    dfs = []
    for code in tqdm(codes):
        client.sync_auto_wait(api.stock_Info())
        df = client.sync_query(api.stock_Info(en_prod_code=code), cn_column_names=True)
        dfs.append(df)
    return pd.concat(dfs)


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    # ret = loop.run_until_complete(test_speed_limit())
    # print(ret)
    results = loop.run_until_complete(test_speed_limit_async())
    print(results)
    # test_speed_limit()
